// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once
#include <limits>
#include <unordered_set>
#include <vector>

namespace kumo::pollux::exec {
    class Driver;
    class JoinBridge;
    class LocalExchangeMemoryManager;
    class MergeSource;
    class MergeJoinSource;
    struct Split;

    /// Corresponds to Presto TaskState, needed for reporting query completion.
    enum class TaskState : int {
        kRunning = 0,
        kFinished = 1,
        kCanceled = 2,
        kAborted = 3,
        kFailed = 4
    };

    std::string taskStateString(TaskState state);

    MELON_ALWAYS_INLINE std::ostream &operator<<(
        std::ostream &os,
        TaskState state) {
        os << taskStateString(state);
        return os;
    }

    struct BarrierState {
        int32_t numRequested;
        std::vector<std::shared_ptr<Driver> > drivers;
        /// Promises given to non-last peer drivers that the last driver will collect
        /// all hashtables from the peers and assembles them into one (HashBuilder
        /// operator does that). After the last drier done its work, the promises are
        /// fulfilled and the non-last drivers can continue.
        std::vector<ContinuePromise> allPeersFinishedPromises;
    };

    /// Structure to accumulate splits for distribution.
    struct SplitsStore {
        /// Arrived (added), but not distributed yet, splits.
        std::deque<exec::Split> splits;
        /// Signal, that no more splits will arrive.
        bool noMoreSplits{false};
        /// Blocking promises given out when out of splits to distribute.
        std::vector<ContinuePromise> splitPromises;
    };

    /// Structure contains the current info on splits for a particular plan node.
    struct SplitsState {
        /// True if the source node is a table scan.
        bool sourceIsTableScan{false};

        /// Plan node-wide 'no more splits'.
        bool noMoreSplits{false};

        /// Keep the max added split's sequence id to deduplicate incoming splits.
        long maxSequenceId{std::numeric_limits<long>::min()};

        /// Map split group id -> split store.
        std::unordered_map<uint32_t, SplitsStore> groupSplitsStores;

        /// We need these due to having promises in the structure.
        SplitsState() = default;

        SplitsState(SplitsState const &) = delete;

        SplitsState &operator=(SplitsState const &) = delete;
    };

    /// Stores local exchange queues with the memory manager.
    struct LocalExchangeState {
        std::shared_ptr<LocalExchangeMemoryManager> memoryManager;
        std::shared_ptr<LocalExchangeVectorPool> vectorPool;
        std::vector<std::shared_ptr<LocalExchangeQueue> > queues;
        std::shared_ptr<common::SkewedPartitionRebalancer>
        scaleWriterPartitionBalancer;
    };

    /// Stores inter-operator state (exchange, bridges) for split groups.
    struct SplitGroupState {
        /// Map from the plan node id of the join to the corresponding JoinBridge.
        /// This map will contain only HashJoinBridge and NestedLoopJoinBridge.
        melon::F14FastMap<core::PlanNodeId, std::shared_ptr<JoinBridge> > bridges;
        /// This map will contain all other custom bridges.
        melon::F14FastMap<core::PlanNodeId, std::shared_ptr<JoinBridge> >
        customBridges;
        /// Holds states for Task::allPeersFinished.
        melon::F14FastMap<core::PlanNodeId, BarrierState> barriers;

        /// Map of merge sources keyed on LocalMergeNode plan node ID.
        std::
        unordered_map<core::PlanNodeId, std::vector<std::shared_ptr<MergeSource> > >
        localMergeSources;

        /// Map of merge join sources keyed on MergeJoinNode plan node ID.
        melon::F14FastMap<core::PlanNodeId, std::shared_ptr<MergeJoinSource> >
        mergeJoinSources;

        /// Map of local exchanges keyed on LocalPartition plan node ID.
        melon::F14FastMap<core::PlanNodeId, LocalExchangeState> localExchanges;

        /// Map of scaled scan controllers keyed on TableScan plan node ID.
        melon::F14FastMap<core::PlanNodeId, std::shared_ptr<ScaledScanController> >
        scaledScanControllers;

        /// Drivers created and still running for this split group.
        /// The split group is finished when this numbers reaches zero.
        uint32_t numRunningDrivers{0};

        /// The number of completed drivers in the output pipeline. When all drivers
        /// in the output pipeline finish, the remaining running pipelines should stop
        /// processing and transition to finished state as well. This happens when
        /// there is a downstream operator that finishes before receiving all input,
        /// e.g. Limit.
        uint32_t numFinishedOutputDrivers{0};

        // True if the state contains structures used for connecting ungrouped
        // execution pipeline with grouped execution pipeline. In that case we don't
        // want to clean up some of these structures.
        bool mixedExecutionMode{false};

        /// Clears the state.
        void clear() {
            if (!mixedExecutionMode) {
                bridges.clear();
                customBridges.clear();
                barriers.clear();
            }
            localMergeSources.clear();
            mergeJoinSources.clear();
            localExchanges.clear();
        }
    };
} // namespace kumo::pollux::exec

template<>
struct fmt::formatter<kumo::pollux::exec::TaskState>
        : formatter<std::string> {
    auto format(kumo::pollux::exec::TaskState state, format_context &ctx)
    const {
        return formatter<std::string>::format(
            kumo::pollux::exec::taskStateString(state), ctx);
    }
};
